#include <vector>
#include <list>
#include <map>
#include <set>
#include <queue>
#include <deque>
#include <stack>
#include <bitset>
#include <algorithm>
#include <functional>
#include <numeric>
#include <utility>
#include <sstream>
#include <iostream>
#include <iomanip>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <ctime>
#include <fstream>
#include <set>

using namespace std;
struct ListNode {
    int val;
    ListNode *next;
    ListNode(int x) : val(x), next(NULL) {}
};

class Solution {
public:
    ListNode *mergeKLists(vector<ListNode *> &lists) {
        ListNode *head = NULL, *ans = NULL;
        priority_queue<pair<int, int>, vector<pair<int, int> >, greater<pair<int, int> > > myQue;
        for (int i = 0; i < lists.size(); i++) {
            if (lists[i] != NULL) {
                myQue.push(make_pair(lists[i]->val, i));
                lists[i] = lists[i]->next;
            }
        }
        while (!myQue.empty()) {
            pair<int, int> a = myQue.top();
            myQue.pop();
            if (!ans )ans = new ListNode(a.first);
			else {ans->next = new ListNode(a.first); ans = ans->next;}
            if (!head) head = ans;

            if (lists[a.second]) {
                myQue.push(make_pair(lists[a.second]->val, a.second));
                lists[a.second] = lists[a.second]->next;
            } 
        }
		return head;
   }
};

//add better code
class comp {
public:
    bool operator()(ListNode* &lhs, ListNode* &rhs)
    {
        return lhs->val >= rhs->val;
    }
};

class Solution {
public:
    ListNode *mergeKLists(vector<ListNode *> &lists) {
        ListNode *cur = new ListNode(0), *head = cur;
        priority_queue<ListNode*, vector<ListNode *>, comp> myPQ;
        vector<int> index(lists.size(), 0);
        for (int i = 0; i < lists.size(); i++) {
            if (lists[i])
                myPQ.push(lists[i]);
        }

        while (!myPQ.empty()) {
            ListNode *top = myPQ.top();
            myPQ.pop();
            cur->next = top;
            cur = cur->next;
            if (top->next)
                myPQ.push(top->next);
        }

        return head->next;
    }
};



ListNode *newList(int data[], int index, int size)
{
    ListNode *node = NULL;
    if (index == size)
        return node;

    node = new ListNode(data[index]);
    node->next = newList(data, index+1, size);

    return node;
}


//test on 1337 online judge and pass all the test
int main(int argc, char **argv)
{
    Solution mySolution;
    cout << "-----------------Test 1--------------------" << endl;
    vector<ListNode *> myList;
    int data[] = {1,4,5,6,7};
    int data2[] = {2,4,5,7,8};
    int data3[] = {2,5,7,8,9};
    ListNode *list;
    list = newList(data, 0, sizeof(data)/sizeof(int));
    myList.push_back(list);
    list = newList(data2, 0, sizeof(data2)/sizeof(int));
    myList.push_back(list);
    list = newList(data3, 0, sizeof(data3)/sizeof(int));
    myList.push_back(list);
	list = mySolution.mergeKLists(myList);
    while(list) {
        cout << list->val << " ";
        list= list->next;
    }
    cout << endl;
    cout << "-----------------Test 2--------------------" << endl;


    cout << "-----------------Test 3--------------------" << endl;


    cout << "-----------------Test 4--------------------" << endl;


    cout << "-----------------Test 5--------------------" << endl;



}
